import asyncio
import json
import logging
import threading
from collections import deque
from datetime import datetime
from typing import Dict, Any, List
import httpx
import websockets
from fastmcp import FastMCP

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 初始化FastMCP实例
mcp = FastMCP("OptimizedLiveStreamBridge")

# 全局变量
websocket_server = None
server_running = False
clients = set()  # 存储所有连接的客户端，限制不超过3个
data_queue = deque()  # 用于存储待处理的数据队列
queue_lock = threading.Lock()  # 队列锁
processing_task = None  # 数据处理任务

# HTTPS目标服务器地址（可修改）
HTTPS_TARGET_URL = "https://example.com/api/live-stream"

# 直播间消息类型定义
MESSAGE_TYPES = {
    "WebcastChatMessage": "聊天消息",
    "WebcastMemberMessage": "用户进入直播间消息",
    "WebcastRoomStatsMessage": "房间状态消息",
    "WebcastRoomRankMessage": "直播间排名消息",
    "WebcastRoomUserSeqMessage": "网络直播间用户序列消息",
    "WebcastLikeMessage": "点赞消息"
}

def classify_message(message: Dict[str, Any]) -> str:
    """
    对消息进行分类
    
    Args:
        message: 消息数据
        
    Returns:
        消息类型描述
    """
    method = message.get("method", "Unknown")
    return MESSAGE_TYPES.get(method, f"未知类型: {method}")

def restructure_websocket_data_for_https(raw_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    重构WebSocket接收的数据以满足向HTTPS服务器转发的格式
    
    Args:
        raw_data: WebSocket接收的原始数据
        
    Returns:
        重构后的数据格式，符合要求的转发格式
    """
    # 根据要求，构造转发到HTTPS服务器的数据结构
    structured_data = {
        "bot_id": "756554292238136XXXX",  # 固定bot_id
        "workflow_id": "756730868947428XXXX",  # 固定workflow_id
        "parameters": raw_data  # 原始数据作为parameters
    }
    
    return structured_data

async def forward_to_https(data: Dict[str, Any]) -> Dict[str, Any]:
    """
    将数据转发到HTTPS服务器并返回响应
    
    Args:
        data: 要转发的数据
        
    Returns:
        来自HTTPS服务器的响应数据
    """
    async with httpx.AsyncClient() as client:
        try:
            # 重构数据格式以满足HTTPS服务器要求
            structured_data = restructure_websocket_data_for_https(data)
            
            # 发送POST请求到HTTPS服务器
            response = await client.post(
                HTTPS_TARGET_URL,
                json=structured_data,
                timeout=30.0
            )
            response.raise_for_status()
            
            # 解析响应为JSON
            result = response.json()
            logger.info(f"HTTPS转发成功: {result}")
            return result
            
        except Exception as e:
            logger.error(f"HTTPS转发失败: {e}")
            return {"error": str(e), "status": "failed"}

async def process_data_queue():
    """
    处理数据队列中的数据，逐条转发到HTTPS服务器
    """
    global data_queue, queue_lock
    
    while True:
        try:
            # 检查队列是否有数据
            with queue_lock:
                if not data_queue:
                    await asyncio.sleep(0.1)  # 短暂休眠避免过度占用CPU
                    continue
            
            # 获取队列中的第一条数据
            with queue_lock:
                if data_queue:
                    data_item = data_queue.popleft()
                else:
                    continue
            
            # 记录消息类型
            message_type = classify_message(data_item)
            logger.info(f"处理消息类型: {message_type}")
            
            # 转发数据到HTTPS服务器
            response_data = await forward_to_https(data_item)
            
            # 将响应数据发送给大模型（通过MCP工具）
            await mcp.notify("live_stream_response", {
                "timestamp": datetime.now().isoformat(),
                "message_type": message_type,
                "original_data": data_item,
                "https_response": response_data
            })
            
            # 短暂休眠避免请求过于频繁
            await asyncio.sleep(0.01)
            
        except Exception as e:
            logger.error(f"处理数据队列时出错: {e}")
            await asyncio.sleep(1)  # 出错时休眠更长时间

async def handle_websocket_client(websocket: websockets.WebSocketServerProtocol, path: str):
    """
    处理WebSocket客户端连接
    
    Args:
        websocket: WebSocket连接对象
        path: 请求路径
    """
    global clients, data_queue, queue_lock
    
    # 检查连接数是否超过限制
    if len(clients) >= 3:
        logger.warning("连接数已达上限，拒绝新连接")
        await websocket.send(json.dumps({
            "error": "服务器连接数已达上限（3个），请稍后再试"
        }, ensure_ascii=False))
        await websocket.close(code=1000, reason="连接数超限")
        return
    
    # 添加客户端到连接集合
    clients.add(websocket)
    client_id = id(websocket)
    logger.info(f"新的直播间客户端连接建立，客户端ID: {client_id}，当前连接数: {len(clients)}")
    
    try:
        # 通知客户端连接已建立
        await websocket.send(json.dumps({
            "status": "connected",
            "message": "直播间WebSocket连接已建立",
            "server": "OptimizedLiveStreamBridge",
            "clientId": client_id
        }, ensure_ascii=False))
        
        # 持续监听客户端消息
        async for message in websocket:
            try:
                # 解析接收到的JSON数据
                raw_data = json.loads(message)
                logger.info(f"接收到直播间客户端数据: {raw_data}")
                
                # 处理数据（可能是单条或批量数据）
                if isinstance(raw_data, list):
                    # 批量数据，逐条添加到队列
                    with queue_lock:
                        for item in raw_data:
                            data_queue.append(item)
                    logger.info(f"已将{len(raw_data)}条数据添加到处理队列")
                else:
                    # 单条数据，添加到队列
                    with queue_lock:
                        data_queue.append(raw_data)
                    logger.info("已将1条数据添加到处理队列")
                
                # 确认接收数据
                await websocket.send(json.dumps({
                    "status": "received",
                    "message": "数据已接收并加入处理队列",
                    "count": len(raw_data) if isinstance(raw_data, list) else 1
                }, ensure_ascii=False))
                
            except json.JSONDecodeError:
                error_msg = {"error": "无效的JSON格式"}
                await websocket.send(json.dumps(error_msg, ensure_ascii=False))
                
            except Exception as e:
                logger.error(f"处理直播间客户端消息时出错: {e}")
                error_msg = {"error": f"处理消息时发生错误: {str(e)}"}
                await websocket.send(json.dumps(error_msg, ensure_ascii=False))
                
    except websockets.exceptions.ConnectionClosed:
        logger.info(f"直播间客户端连接已关闭，客户端ID: {client_id}")
    except Exception as e:
        logger.error(f"处理直播间客户端连接时出错: {e}")
    finally:
        # 清理客户端连接
        clients.discard(websocket)
        logger.info(f"直播间客户端连接处理完成，客户端ID: {client_id}，当前连接数: {len(clients)}")

@mcp.tool("send_data", description="向直播间WebSocket客户端发送数据")
async def send_data(data: Dict[str, Any]) -> bool:
    """
    向当前连接的直播间WebSocket客户端发送数据
    
    Args:
        data: 要发送的数据
        
    Returns:
        发送是否成功
    """
    global clients
    
    if not clients:
        logger.warning("当前无直播间客户端连接")
        return False
    
    success_count = 0
    # 向所有连接的客户端发送数据
    for client in list(clients):  # 创建副本以防在迭代时修改集合
        try:
            await client.send(json.dumps(data, ensure_ascii=False))
            success_count += 1
        except Exception as e:
            logger.error(f"向客户端发送数据失败: {e}")
            # 移除已断开的连接
            clients.discard(client)
    
    return success_count > 0

@mcp.tool("get_connection_status", description="获取直播间WebSocket连接状态")
def get_connection_status() -> Dict[str, Any]:
    """
    获取当前直播间WebSocket连接状态
    
    Returns:
        包含连接状态信息的字典
    """
    global clients, data_queue
    
    with queue_lock:
        queue_size = len(data_queue)
    
    return {
        "connected_clients": len(clients),
        "max_connections": 3,
        "queue_size": queue_size,
        "clients_info": [id(client) for client in clients]
    }

@mcp.tool("get_message_types", description="获取支持的消息类型")
def get_message_types() -> Dict[str, Any]:
    """
    获取支持的消息类型及其描述
    
    Returns:
        包含消息类型信息的字典
    """
    return {
        "supported_types": MESSAGE_TYPES,
        "total_types": len(MESSAGE_TYPES)
    }

@mcp.tool("start_live_stream_server", description="启动直播间服务")
async def start_live_stream_server() -> Dict[str, Any]:
    """
    启动直播间WebSocket服务器
    
    Returns:
        启动结果信息
    """
    global websocket_server, server_running, processing_task
    
    if server_running:
        return {
            "status": "already_running",
            "message": "直播间服务器已在运行中",
            "address": "ws://localhost:4433"
        }
    
    try:
        # 启动WebSocket服务器
        websocket_server = await websockets.serve(
            handle_websocket_client,
            "localhost",
            4433,
            ping_interval=None  # 禁用自动ping以减少干扰
        )
        
        # 启动数据处理任务
        processing_task = asyncio.create_task(process_data_queue())
        
        server_running = True
        logger.info("直播间WebSocket服务器已在 ws://localhost:4433 启动并开始监听...")
        
        return {
            "status": "success",
            "message": "直播间服务器启动成功",
            "address": "ws://localhost:4433"
        }
    except Exception as e:
        logger.error(f"启动直播间服务器失败: {e}")
        return {
            "status": "error",
            "message": f"启动直播间服务器失败: {str(e)}"
        }

@mcp.tool("stop_live_stream_server", description="停止直播间服务")
async def stop_live_stream_server() -> Dict[str, Any]:
    """
    停止直播间WebSocket服务器
    
    Returns:
        停止结果信息
    """
    global websocket_server, server_running, processing_task, clients, data_queue
    
    if not server_running:
        return {
            "status": "not_running",
            "message": "直播间服务器未在运行"
        }
    
    try:
        # 关闭所有客户端连接
        for client in list(clients):
            try:
                await client.close()
            except:
                pass
        
        # 清理连接集合和数据队列
        clients.clear()
        with queue_lock:
            data_queue.clear()
        
        # 关闭WebSocket服务器
        if websocket_server:
            websocket_server.close()
            await websocket_server.wait_closed()
        
        # 取消数据处理任务
        if processing_task and not processing_task.done():
            processing_task.cancel()
            try:
                await processing_task
            except asyncio.CancelledError:
                pass
        
        server_running = False
        logger.info("直播间WebSocket服务器已关闭")
        
        return {
            "status": "success",
            "message": "直播间服务器已停止"
        }
    except Exception as e:
        logger.error(f"停止直播间服务器失败: {e}")
        return {
            "status": "error",
            "message": f"停止直播间服务器失败: {str(e)}"
        }

@mcp.prompt("连接直播间")
async def connect_live_stream():
    """
    当大模型输入"连接直播间"时触发此函数
    启动MCP服务端并开始监听、转发
    """
    logger.info("收到'连接直播间'指令，正在启动直播间服务...")
    
    # 启动直播间服务器
    result = await start_live_stream_server()
    
    if result["status"] == "success":
        # 启动FastMCP服务
        logger.info("直播间服务启动成功，正在启动MCP服务...")
        return {
            "status": "success",
            "message": "直播间连接已建立，MCP服务已启动",
            "server_address": result["address"]
        }
    else:
        return {
            "status": "error",
            "message": f"直播间服务启动失败: {result['message']}"
        }

async def main():
    """
    主函数：启动FastMCP服务
    """
    logger.info("正在启动优化版直播间MCP服务...")
    await mcp.run()

if __name__ == "__main__":
    asyncio.run(main())